import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

prompt_template = ChatPromptTemplate.from_template(
    "请写一首关于{topic}的诗，要求不超过{num}行。"
    )
chat_model = ChatZhipuAI(model=model,
                        temperature=0.9,              
    )
output_parser = StrOutputParser()

# prompt = prompt_template.format(topic="春天", num=8)
# response = chat_model.invoke(prompt)
# result= output_parser.invoke(response)
chain = prompt_template | chat_model | output_parser
result = chain.invoke({"topic":"春天", "num":8})
print(result)
